# -*- coding: utf-8 -*-
#!/usr/bin/python


import os
import sys
cur_path =os.path.abspath(os.path.join(os.path.dirname('__file__'), os.path.pardir))
sys.path.append(cur_path)

from psm_conf import *
import socket
gLocalIp=socket.gethostbyname(socket.getfqdn(socket.gethostname()))

gWorkName=gLocalIp+":"+str(PORT_PSM_WORKER)

import time


import json
from worker_src import ps_moniter_worker
from worker_src.ttypes import *

import thrift
from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol
from thrift.server import TServer
from kazoo.client import KazooClient
from hulk_proxy.HulkProxyServ import Client
from scan_data.ttypes import ScanData

def log(msg):
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())+"|"+msg)

def mygetstatusoutput(cmd):
    """Return (status, output) of executing cmd in a shell."""
    import os
    #pipe = os.popen('{ ' + cmd + ' } 2>&1', 'r')
    #pipe = os.popen('{ ' + cmd + '; } 2>&1', 'r')
    pipe = os.popen(cmd, 'r')
    text = pipe.read()
    sts = pipe.close()
    if sts is None: sts = 0
    if text[-1:] == '\n': text = text[:-1]
    return sts, text

def getMyInfo():
    info={}
    ret, retString = mygetstatusoutput("free -h")
    free_data=retString.split("\n")[1].split()
    free_data.pop(0)
    free_key=["total", "used", "free", "shared", "buff/cache", "available"]
    index=0
    free_ret=""
    while index < len(free_data):
        free_ret+=free_key[index]
        free_ret+=":"
        free_ret+=free_data[index]
        free_ret+="|"
        index+=1
    info["mem"]=free_ret

    ret, retString = mygetstatusoutput("uptime")
    info["load"]="load average: "+retString.split('load average: ')[1]

    ret, retString = mygetstatusoutput("iostat")
    iostat_key=retString.split("\n")[2].split()
    iostat_data=retString.split("\n")[3].split()
    iostat_ret=""
    iostat_ret_ket=""
    iostat_ret_ket+=iostat_key[0]
    iostat_ret_ket+="(iostat)"
    iostat_key.pop(0)
    index=0
    while index < len(iostat_key):
        iostat_ret+=iostat_key[index]
        iostat_ret+=":"
        iostat_ret+=iostat_data[index]
        iostat_ret+="|"
        index+=1
    info["cpu"]=iostat_ret

    ret, retString = mygetstatusoutput("df -h /data |grep -v Filesystem")
    df_key=["Size", "Used", "Avail", "Use%", "Mounted on"]
    df_data=retString.split()
    df_data.pop(0)
    df_ret=""
    index=0
    while index < len(df_key):
        df_ret+=df_key[index]
        df_ret+=":"
        df_ret+=df_data[index]
        df_ret+="|"
        index+=1
    info["disk"]=df_ret

    ret, retString = mygetstatusoutput("ps -ef|grep ./ParaServ")
    info["ParaServ_num"]=len(retString.split("\n"))-2
    return json.dumps(info)


class workerHandler:
    def __init__(self):
        self.step=0
        pass

    def getPort(self):
        new_port=-1
        tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
        transport = TTransport.TFramedTransport(tsocket)
        protocol = TBinaryProtocol.TBinaryProtocol(transport)
        client = Client(protocol)
        transport.open()
        tmp_hash = "psm_"+gWorkName
        iPort = int(client.hget(tmp_hash, "port"))
        time=0
        # port取值范围是[28000, 29000)
        while time<1000:
            ret, retString = mygetstatusoutput("lsof -i:"+str(iPort))
            if retString:
                log("reuse port. lsof -i:"+str(iPort))
                iPort+=1
            else:
                client.hpush(tmp_hash, "port", str(iPort+1))
                log("use port "+str(iPort))
                new_port = iPort
                break
            time +=1
        #查找可用端口
        transport.close()
        if new_port == -1:
            return "retry"
        return str(new_port)

    def pwd(self):
        log(str(self.step))
        self.step = self.step+1 if self.step+1 <99999 else 0
        return "aa"
    
    def stopPs(self, pid_modelNameId):
        log("stopPs|in|"+pid_modelNameId)
        pid=pid_modelNameId.split("|")[0]
        model_name_id=pid_modelNameId.split("|")[1]
         
        log("rm -rf /data/rocksdb/"+model_name_id)
        log("rm -rf /data/transport/"+model_name_id)
        ret, retString = mygetstatusoutput("rm -rf /data/rocksdb/"+model_name_id)
        ret, retString = mygetstatusoutput("rm -rf /data/transport/"+model_name_id)

        ret, retString = mygetstatusoutput("ps -ef|grep "+str(pid))
        if len(retString.split("\n")) == 3:
            mygetstatusoutput("kill -9 "+str(pid))

        if pid_modelNameId.find("proxy") == -1:
            # 不是proxy, 则是ParaServ, 更新本机ParaServ进程记录数
            tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
            transport = TTransport.TFramedTransport(tsocket)
            protocol = TBinaryProtocol.TBinaryProtocol(transport)
            client = Client(protocol)
            transport.open()
            tmp_hash = "psm_"+gWorkName
            ret = client.hget(tmp_hash, "ps_number")
            log("ParaServ -1, now running %d." %(int(ret)-1))
            client.hmpush(tmp_hash, {"ps_number":str(int(ret)-1)})
            transport.close()

        log("stopPs|out|"+pid_modelNameId)
        return str(ret)
      
    def runPs(self, input):
        log("runPs|in|"+input)
        j = json.loads(input)
        model_name_id = j["model_name"]+"_"+str(j["model_id"])
        machine_type=j["machine_type"]
        groupid=j["groupid"]

        tp=self.getPort()
        if tp == "retry":
            return {-1:"no -port"}

        mp = self.getPort()
        if mp == "retry":
            return {-1:"no -port"}
        cmd = 'nohup ./ParaServ -g {groupid} -thread_num {thread_num} -mn {model_name} -tp {thrift_port} -mp {master_port} -mt {machine_type} > psServ_{model_name}.log  2>&1 &'.format(
                           groupid=groupid, thread_num=ParaServ_THREAD_NUMBER, model_name= model_name_id, thrift_port=tp, master_port=mp, machine_type=machine_type)   
        log(cmd)
        ret, retString = mygetstatusoutput(cmd) 
        if ret is not 0:
            return {ret:retString}

        sleep_times=0
        check_time=10
        while sleep_times < check_time:
            log(str(sleep_times))
            time.sleep(0.01)# 每次sleep 0.01秒
            ret, retString = mygetstatusoutput("lsof -i:"+str(tp))
            if len(retString.split("\n")) >= 2 and retString.split("\n")[1].split()[0] == "ParaServ":
                break
            elif sleep_times+1>=check_time: # 超过次数直接返回
                return {-3:"error: lsof -i:"+str(tp)}
            sleep_times+=1

        # 更新本机ParaServ进程记录数
        tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
        transport = TTransport.TFramedTransport(tsocket)
        protocol = TBinaryProtocol.TBinaryProtocol(transport)
        client = Client(protocol)
        transport.open()
        tmp_hash = "psm_"+gWorkName
        ret = client.hget(tmp_hash, "ps_number")
        log("ParaServ +1, now running %d." %(int(ret)+1))
        client.hmpush(tmp_hash, {"ps_number":str(int(ret)+1)})
        transport.close()
        # 返回pid|groupid|master_port
        return {0:retString.split("\n")[1].split()[1]+"|"+groupid+"|"+str(mp)}

    def runPsProxy(self, input):
        j = json.loads(input)
        model_name_id = j["model_name"]+"_"+str(j["model_id"])
        machine_type=j["machine_type"]
        origin_psProxy_zkUrl=""
        origin_psProxy_zkPath=""
        if machine_type == "realtime":
            origin_psProxy_zkUrl=ZK_URL
            origin_psProxy_zkPath="{NEW_ZKPATH}/psProxy/{MODEL_NAME_ID}_{MACHINE_TYPE}".format( NEW_ZKPATH=ZK_PATH, MODEL_NAME_ID=model_name_id, MACHINE_TYPE="offline")
        elif machine_type == "online":
            origin_psProxy_zkUrl=ZK_URL
            origin_psProxy_zkPath="{NEW_ZKPATH}/psProxy/{MODEL_NAME_ID}_{MACHINE_TYPE}".format( NEW_ZKPATH=ZK_PATH, MODEL_NAME_ID=model_name_id, MACHINE_TYPE="realtime")
        
        psProxy_port=self.getPort()
        if psProxy_port == "retry":
            return {-1:"no -port"}
            
        #创建tcmcp永久结点，stop时删除 
        tcmcp_path="{NEW_ZKPATH}/psProxy/{MODEL_NAME_ID}_{MACHINE_TYPE}".format(NEW_ZKPATH=ZK_PATH, MODEL_NAME_ID=model_name_id, MACHINE_TYPE=machine_type)
        log("create tcmcp_path :"+tcmcp_path)           
        zkc = KazooClient(hosts=ZK_URL, timeout=100)
        zkc.start()
        zkc.create(tcmcp_path, (model_name_id+"_"+machine_type).encode(encoding="utf-8"))

        cmd='nohup ./PSProxy  {NEW_ZKURL} {TCMCP_PATH} {PORT} {MODEL_NAME_ID_MT} {RUNPSNUM} {ORIGIN_PSPROXY_ZKURL} {ORIGIN_PSPROXY_ZKPATH} >psProxy_{MODEL_NAME_ID}_{MACHINE_TYPE}>.log  2>&1 &'.format(
                        MACHINE_TYPE=machine_type, NEW_ZKURL=ZK_URL, TCMCP_PATH=tcmcp_path, MODEL_NAME_ID=model_name_id, MODEL_NAME_ID_MT=model_name_id+"_"+machine_type, PORT=psProxy_port, ORIGIN_PSPROXY_ZKURL=origin_psProxy_zkUrl, ORIGIN_PSPROXY_ZKPATH=origin_psProxy_zkPath, RUNPSNUM=RUN_ParaServ_NUMBER[machine_type], IP=gLocalIp)
        #cmd="nohup ./PSProxy %s %s %s/psProxy/%s/psm_psProxy %s %s >psProxy_%s.log  2>&1 &" % (machine_type, ZK_URL, ZK_PATH, model_name_id, port, model_name_id, model_name_id)
        ret, retString = mygetstatusoutput(cmd) 
        log(cmd)
        if ret is not 0:
            return {ret:retString}
        time.sleep(1)
        sleep_times=0
        check_time=100
        while sleep_times < check_time:
            log("check PSProxy ("+str(sleep_times)+"/"+str(check_time)+") times")
            time.sleep(0.01)# 每次sleep 0.01秒
            ret, retString = mygetstatusoutput("lsof -i:"+str(psProxy_port))
            if len(retString.split("\n")) >= 2 and retString.split("\n")[1].split()[0] == "PSProxy":
                break
            elif sleep_times+1>=check_time: # 超过次数直接返回
                return {-3:"error: lsof -i:"+str(psProxy_port)}
            sleep_times+=1
        lsof=retString.split("\n")
        if len(lsof) is not 2:
            return {-3:"error: lsof -i:"+str(psProxy_port)}
        
        # 返回本次启动进程的pid
        return {0:lsof[1].split()[1]}

    def getInfo(self, input):
        log("getInfo")
        return getMyInfo()


def registe2zk():
    zkc = KazooClient(hosts=ZK_URL, timeout=100)
    zkc.start()
    # 判断节点是否存在
    if zkc.exists(ZK_PATH+"/"+ MY_MACHINE_TYPE):
        log(ZK_PATH+"/"+ MY_MACHINE_TYPE + "存在")
    else:
        # 建立节点，成功后返回新节点路径
        childrenPath = zkc.create(ZK_PATH+"/"+ MY_MACHINE_TYPE,  MY_MACHINE_TYPE.encode(encoding="utf-8"))
        log("创建节点：" + childrenPath + "成功。")

    # 创建临时节点，连接断开则节点自动删除
    zkc.create(ZK_PATH+"/"+ MY_MACHINE_TYPE+"/"+gWorkName, gWorkName.encode(encoding="utf-8"), ephemeral=True)



def initHulk():
    tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
    transport = TTransport.TFramedTransport(tsocket)
    protocol = TBinaryProtocol.TBinaryProtocol(transport)
    client = Client(protocol)
    transport.open()
    tmp_hash = "psm_"+gWorkName
    ret = client.hget(tmp_hash, "ps_number")
    if not ret:
        client.hmpush(tmp_hash, {"ps_number":"0", "port":"28000"})
        log("创建 " + tmp_hash)
    else:
        log(tmp_hash+" 已存在, "+"重置ps_number")
        ret, retString = mygetstatusoutput("ps -ef|grep ParaServ")
        client.hmpush(tmp_hash, {"ps_number":str(len(retString.split('\n'))-2)})
        ret = client.hget(tmp_hash, "ps_number")
    transport.close()


if __name__=="__main__":
    registe2zk()

    initHulk()

    handler = workerHandler()
    processor = ps_moniter_worker.Processor(handler)
    transport = TSocket.TServerSocket(gLocalIp, PORT_PSM_WORKER)
    tfactory = TTransport.TBufferedTransportFactory()
    pfactory = TBinaryProtocol.TBinaryProtocolFactory()
    server = TServer.TSimpleServer(processor, transport, tfactory, pfactory)
    log("Starting python server...")
    server.serve()
